import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.feature_selection import SelectKBest, f_regression
from sklearn.impute import SimpleImputer
import warnings
warnings.filterwarnings('ignore')

class AmesFeatureEngineer:
    def __init__(self):
        self.scaler = StandardScaler()
        self.label_encoders = {}
        self.feature_selector = None
        self.selected_features = None
        self.imputer = SimpleImputer(strategy='median')
        
    def create_new_features(self, df):
        """
        创建新的组合特征
        """
        df = df.copy()
        
        # 1. 面积相关特征
        if all(col in df.columns for col in ['TotalBsmtSF', '1stFlrSF', '2ndFlrSF']):
            df['TotalSF'] = df['TotalBsmtSF'] + df['1stFlrSF'] + df['2ndFlrSF']
        
        if all(col in df.columns for col in ['TotalSF', 'GarageArea', 'WoodDeckSF', 'OpenPorchSF']):
            df['TotalArea'] = df['TotalSF'] + df['GarageArea'] + df['WoodDeckSF'] + df['OpenPorchSF']
        
        if all(col in df.columns for col in ['FullBath', 'HalfBath', 'BsmtFullBath', 'BsmtHalfBath']):
            df['TotalBath'] = df['FullBath'] + 0.5 * df['HalfBath'] + df['BsmtFullBath'] + 0.5 * df['BsmtHalfBath']
        
        if all(col in df.columns for col in ['OpenPorchSF', 'EnclosedPorch', '3SsnPorch', 'ScreenPorch']):
            df['TotalPorch'] = df['OpenPorchSF'] + df['EnclosedPorch'] + df['3SsnPorch'] + df['ScreenPorch']
        
        # 2. 比率特征
        if all(col in df.columns for col in ['GrLivArea', 'LotArea']):
            df['LivingAreaRatio'] = df['GrLivArea'] / (df['LotArea'].replace(0, 1) + 1)  # 避免除零
        
        if all(col in df.columns for col in ['TotalBsmtSF', 'TotalSF']):
            df['BsmtRatio'] = df['TotalBsmtSF'] / (df['TotalSF'].replace(0, 1) + 1)
        
        if all(col in df.columns for col in ['GarageArea', 'TotalSF']):
            df['GarageAreaRatio'] = df['GarageArea'] / (df['TotalSF'].replace(0, 1) + 1)
        
        # 3. 房屋年龄相关特征
        if all(col in df.columns for col in ['YrSold', 'YearBuilt']):
            df['HouseAge'] = df['YrSold'] - df['YearBuilt']
        
        if all(col in df.columns for col in ['YrSold', 'YearRemodAdd']):
            df['RemodAge'] = df['YrSold'] - df['YearRemodAdd']
            df['IsRemodeled'] = (df['YearBuilt'] != df['YearRemodAdd']).astype(int)
        
        if all(col in df.columns for col in ['YrSold', 'YearBuilt']):
            df['IsNewHouse'] = (df['YrSold'] - df['YearBuilt'] <= 1).astype(int)
        
        # 4. 房间配置特征
        if all(col in df.columns for col in ['TotRmsAbvGrd', 'GrLivArea']):
            df['RoomDensity'] = df['TotRmsAbvGrd'] / (df['GrLivArea'].replace(0, 1) + 1)
        
        if all(col in df.columns for col in ['BedroomAbvGr', 'TotRmsAbvGrd']):
            df['BedroomRatio'] = df['BedroomAbvGr'] / (df['TotRmsAbvGrd'].replace(0, 1) + 1)
        
        return df
    
    def encode_categorical_features(self, df, is_training=True):
        """
        对分类特征进行编码
        """
        df = df.copy()
        
        # 有序分类特征的映射
        quality_mapping = {'Ex': 5, 'Gd': 4, 'TA': 3, 'Fa': 2, 'Po': 1, np.nan: 0, 'NA': 0, 'None': 0}
        exposure_mapping = {'Gd': 4, 'Av': 3, 'Mn': 2, 'No': 1, np.nan: 0, 'NA': 0, 'None': 0}
        finish_mapping = {'GLQ': 6, 'ALQ': 5, 'BLQ': 4, 'Rec': 3, 'LwQ': 2, 'Unf': 1, np.nan: 0, 'NA': 0, 'None': 0}
        functional_mapping = {'Typ': 7, 'Min1': 6, 'Min2': 5, 'Mod': 4, 'Maj1': 3, 'Maj2': 2, 'Sev': 1, 'Sal': 0}
        pool_mapping = {'Ex': 4, 'Gd': 3, 'TA': 2, 'Fa': 1, np.nan: 0, 'NA': 0, 'None': 0}
        garage_finish_mapping = {'Fin': 3, 'RFn': 2, 'Unf': 1, np.nan: 0, 'NA': 0, 'None': 0}
        fence_mapping = {'GdPrv': 4, 'MnPrv': 3, 'GdWo': 2, 'MnWw': 1, np.nan: 0, 'NA': 0, 'None': 0}
        
        # 应用有序编码
        quality_features = ['ExterQual', 'ExterCond', 'BsmtQual', 'BsmtCond', 
                          'HeatingQC', 'KitchenQual', 'FireplaceQu', 'GarageQual', 'GarageCond']
        
        for feature in quality_features:
            if feature in df.columns:
                df[f'{feature}_encoded'] = df[feature].map(quality_mapping).fillna(0)
        
        # PoolQC特殊处理
        if 'PoolQC' in df.columns:
            df['PoolQC_encoded'] = df['PoolQC'].map(pool_mapping).fillna(0)
        
        # GarageFinish特殊处理
        if 'GarageFinish' in df.columns:
            df['GarageFinish_encoded'] = df['GarageFinish'].map(garage_finish_mapping).fillna(0)
        
        # Fence特殊处理
        if 'Fence' in df.columns:
            df['Fence_encoded'] = df['Fence'].map(fence_mapping).fillna(0)
        
        # 其他有序特征
        if 'BsmtExposure' in df.columns:
            df['BsmtExposure_encoded'] = df['BsmtExposure'].map(exposure_mapping).fillna(0)
        
        if 'BsmtFinType1' in df.columns:
            df['BsmtFinType1_encoded'] = df['BsmtFinType1'].map(finish_mapping).fillna(0)
        
        if 'BsmtFinType2' in df.columns:
            df['BsmtFinType2_encoded'] = df['BsmtFinType2'].map(finish_mapping).fillna(0)
        
        if 'Functional' in df.columns:
            df['Functional_encoded'] = df['Functional'].map(functional_mapping).fillna(7)
        
        # 对无序分类特征使用标签编码
        categorical_features = ['MSZoning', 'Street', 'Alley', 'LotShape', 'LandContour', 
                               'Utilities', 'LotConfig', 'LandSlope', 'Neighborhood', 
                               'Condition1', 'Condition2', 'BldgType', 'HouseStyle', 
                               'RoofStyle', 'RoofMatl', 'Exterior1st', 'Exterior2nd', 
                               'MasVnrType', 'Foundation', 'Heating', 'CentralAir', 
                               'Electrical', 'GarageType', 'PavedDrive', 'MiscFeature', 
                               'SaleType', 'SaleCondition']
        
        for feature in categorical_features:
            if feature in df.columns:
                # 处理缺失值
                df[feature] = df[feature].fillna('Missing')
                
                if is_training:
                    le = LabelEncoder()
                    df[f'{feature}_encoded'] = le.fit_transform(df[feature].astype(str))
                    self.label_encoders[feature] = le
                else:
                    le = self.label_encoders.get(feature)
                    if le is not None:
                        # 处理测试集中出现的新类别
                        df[f'{feature}_encoded'] = df[feature].astype(str).apply(
                            lambda x: le.transform([x])[0] if x in le.classes_ else -1
                        )
                    else:
                        df[f'{feature}_encoded'] = 0
        
        # 5. 外部特征组合 - 需要在编码后创建
        if all(col in df.columns for col in ['ExterQual_encoded', 'ExterCond_encoded']):
            df['ExteriorScore'] = df['ExterQual_encoded'] * df['ExterCond_encoded']
        
        if all(col in df.columns for col in ['BsmtQual_encoded', 'BsmtCond_encoded']):
            df['BasementScore'] = df['BsmtQual_encoded'] * df['BsmtCond_encoded']
        
        if all(col in df.columns for col in ['GarageQual_encoded', 'GarageCond_encoded']):
            df['GarageScore'] = df['GarageQual_encoded'] * df['GarageCond_encoded']
        
        # 6. 位置和邻里特征增强
        if 'Condition1' in df.columns:
            df['IsNearPositive'] = ((df['Condition1'].isin(['PosN', 'PosA'])) | 
                                   (df['Condition1'].str.contains('Pos', na=False))).astype(int)
        
        if 'Condition1' in df.columns:
            df['IsNearRailroad'] = ((df['Condition1'].str.contains('RR', na=False))).astype(int)
        
        if 'Condition2' in df.columns:
            df['IsNearPositive'] = df['IsNearPositive'] | ((df['Condition2'].isin(['PosN', 'PosA'])) | 
                                   (df['Condition2'].str.contains('Pos', na=False))).astype(int)
            df['IsNearRailroad'] = df['IsNearRailroad'] | (df['Condition2'].str.contains('RR', na=False)).astype(int)
        
        return df
    
    def handle_numerical_features(self, df, is_training=True):
        """
        处理数值特征
        """
        df = df.copy()
        
        # 对偏态分布的数值特征进行对数变换
        skewed_features = ['LotFrontage', 'LotArea', 'MasVnrArea', 'BsmtFinSF1', 
                          'BsmtFinSF2', 'BsmtUnfSF', 'TotalBsmtSF', '1stFlrSF', 
                          '2ndFlrSF', 'GrLivArea', 'GarageArea', 'WoodDeckSF', 
                          'OpenPorchSF', 'EnclosedPorch', '3SsnPorch', 'ScreenPorch', 
                          'PoolArea', 'MiscVal']
        
        for feature in skewed_features:
            if feature in df.columns:
                # 加1避免对0或负值取对数
                df[f'{feature}_log'] = np.log1p(df[feature].fillna(0))
        
        # 标准化数值特征
        numerical_features = ['OverallQual', 'OverallCond', 'YearBuilt', 'YearRemodAdd', 
                             'TotalBsmtSF', '1stFlrSF', '2ndFlrSF', 'GrLivArea', 
                             'FullBath', 'HalfBath', 'BedroomAbvGr', 'KitchenAbvGr', 
                             'TotRmsAbvGrd', 'Fireplaces', 'GarageYrBlt', 'GarageCars', 
                             'MoSold', 'YrSold']
        
        features_to_scale = [f for f in numerical_features if f in df.columns]
        
        if features_to_scale:
            if is_training:
                df[features_to_scale] = self.scaler.fit_transform(df[features_to_scale])
            else:
                df[features_to_scale] = self.scaler.transform(df[features_to_scale])
        
        return df
    
    def ensure_numeric_data(self, df):
        """
        确保所有数据都是数值类型
        """
        df = df.copy()
        
        # 检查非数值列
        non_numeric_cols = df.select_dtypes(include=['object']).columns
        if len(non_numeric_cols) > 0:
            print(f"发现非数值列: {list(non_numeric_cols)}")
            # 对非数值列进行标签编码
            for col in non_numeric_cols:
                df[col] = LabelEncoder().fit_transform(df[col].astype(str))
        
        return df
    
    def handle_missing_values(self, X_train, X_test):
        """
        处理缺失值
        """
        # 检查缺失值
        train_missing = X_train.isnull().sum().sum()
        test_missing = X_test.isnull().sum().sum()
        
        print(f"训练集缺失值数量: {train_missing}")
        print(f"测试集缺失值数量: {test_missing}")
        
        if train_missing > 0 or test_missing > 0:
            # 使用中位数填充缺失值
            self.imputer.fit(X_train)
            X_train_filled = pd.DataFrame(
                self.imputer.transform(X_train),
                columns=X_train.columns,
                index=X_train.index
            )
            X_test_filled = pd.DataFrame(
                self.imputer.transform(X_test),
                columns=X_test.columns,
                index=X_test.index
            )
            print("已使用中位数填充缺失值")
            return X_train_filled, X_test_filled
        
        return X_train, X_test
    
    def feature_selection(self, X_train, y_train, X_test, k=100):
        """
        特征选择
        """
        # 确保数据都是数值类型
        X_train = self.ensure_numeric_data(X_train)
        X_test = self.ensure_numeric_data(X_test)
        
        # 处理缺失值
        X_train, X_test = self.handle_missing_values(X_train, X_test)
        
        self.feature_selector = SelectKBest(score_func=f_regression, k=min(k, X_train.shape[1]))
        X_train_selected = self.feature_selector.fit_transform(X_train, y_train)
        X_test_selected = self.feature_selector.transform(X_test)
        
        # 获取选择的特征名
        self.selected_features = X_train.columns[self.feature_selector.get_support()]
        
        # 将numpy数组转换回DataFrame
        X_train_df = pd.DataFrame(X_train_selected, columns=self.selected_features, index=X_train.index)
        X_test_df = pd.DataFrame(X_test_selected, columns=self.selected_features, index=X_test.index)
        
        return X_train_df, X_test_df
    
    def get_feature_importance(self):
        """
        获取特征重要性评分
        """
        if self.feature_selector is not None and self.selected_features is not None:
            # 使用特征选择器的分数和选择的特征名
            selected_scores = self.feature_selector.scores_[self.feature_selector.get_support()]
            feature_scores = pd.DataFrame({
                'feature': self.selected_features,
                'score': selected_scores
            })
            return feature_scores.sort_values('score', ascending=False)
        return None
    
    def process_data(self, train_df, test_df, target_col='SalePrice', perform_selection=True):
        """
        完整的数据处理流程
        """
        print("Starting feature engineering process...")
        
        # 分离特征和目标变量
        if target_col in train_df.columns:
            y_train = train_df[target_col]
            X_train = train_df.drop(columns=[target_col])
        else:
            X_train = train_df
            y_train = None
        
        X_test = test_df
        
        # 保存Id列用于最终提交
        train_ids = X_train['Id'] if 'Id' in X_train.columns else None
        test_ids = X_test['Id'] if 'Id' in X_test.columns else None
        
        # 移除Id列（不用于建模）
        if 'Id' in X_train.columns:
            X_train = X_train.drop(columns=['Id'])
        if 'Id' in X_test.columns:
            X_test = X_test.drop(columns=['Id'])
        
        # 1. 编码分类特征
        print("Step 1: Encoding categorical features...")
        X_train_encoded = self.encode_categorical_features(X_train, is_training=True)
        X_test_encoded = self.encode_categorical_features(X_test, is_training=False)
        
        # 2. 创建新特征
        print("Step 2: Creating new features...")
        X_train_new = self.create_new_features(X_train_encoded)
        X_test_new = self.create_new_features(X_test_encoded)
        
        # 3. 处理数值特征
        print("Step 3: Processing numerical features...")
        X_train_processed = self.handle_numerical_features(X_train_new, is_training=True)
        X_test_processed = self.handle_numerical_features(X_test_new, is_training=False)
        
        # 4. 移除原始分类特征（保留编码后的版本）
        original_categorical = ['MSZoning', 'Street', 'Alley', 'LotShape', 'LandContour', 
                               'Utilities', 'LotConfig', 'LandSlope', 'Neighborhood', 
                               'Condition1', 'Condition2', 'BldgType', 'HouseStyle', 
                               'RoofStyle', 'RoofMatl', 'Exterior1st', 'Exterior2nd', 
                               'MasVnrType', 'Foundation', 'Heating', 'CentralAir', 
                               'Electrical', 'GarageType', 'PavedDrive', 'MiscFeature', 
                               'SaleType', 'SaleCondition', 'ExterQual', 'ExterCond', 
                               'BsmtQual', 'BsmtCond', 'HeatingQC', 'KitchenQual', 
                               'FireplaceQu', 'GarageQual', 'GarageCond', 'PoolQC',
                               'BsmtExposure', 'BsmtFinType1', 'BsmtFinType2', 'Functional',
                               'GarageFinish', 'Fence']
        
        cols_to_drop = [col for col in original_categorical if col in X_train_processed.columns]
        X_train_processed = X_train_processed.drop(columns=cols_to_drop, errors='ignore')
        X_test_processed = X_test_processed.drop(columns=cols_to_drop, errors='ignore')
        
        # 5. 特征选择（仅在训练时）
        if perform_selection and y_train is not None:
            print("Step 4: Performing feature selection...")
            X_train_final, X_test_final = self.feature_selection(
                X_train_processed, y_train, X_test_processed
            )
        else:
            X_train_final = X_train_processed
            X_test_final = X_test_processed
        
        # 确保训练集和测试集特征一致
        common_features = list(set(X_train_final.columns) & set(X_test_final.columns))
        X_train_final = X_train_final[common_features]
        X_test_final = X_test_final[common_features]
        
        # 添加Id列回数据
        if train_ids is not None:
            X_train_final['Id'] = train_ids.values
        if test_ids is not None:
            X_test_final['Id'] = test_ids.values
        
        print(f"Final training features: {X_train_final.shape[1]}")
        print(f"Final test features: {X_test_final.shape[1]}")
        
        if y_train is not None:
            return X_train_final, X_test_final, y_train
        else:
            return X_train_final, X_test_final

def main():
    """
    主函数：执行特征工程流程
    """
    # 读取清洗后的数据
    try:
        train_clean = pd.read_csv('train_clean.csv')
        test_clean = pd.read_csv('test_clean.csv')
        print("Successfully loaded cleaned data")
        print(f"Training set shape: {train_clean.shape}")
        print(f"Test set shape: {test_clean.shape}")
    except FileNotFoundError:
        print("Cleaned data files not found. Please ensure train_clean.csv and test_clean.csv exist.")
        return
    
    # 初始化特征工程器
    feature_engineer = AmesFeatureEngineer()
    
    # 执行特征工程
    X_train_fe, X_test_fe, y_train = feature_engineer.process_data(
        train_clean, test_clean, target_col='SalePrice', perform_selection=True
    )
    
    # 保存处理后的特征数据
    output_dir = 'data/'
    import os
    os.makedirs(output_dir, exist_ok=True)
    
    # 合并特征和目标变量用于训练
    train_feature = pd.concat([X_train_fe, y_train], axis=1)
    test_feature = X_test_fe
    
    train_feature.to_csv(f'{output_dir}train_feature.csv', index=False)
    test_feature.to_csv(f'{output_dir}test_feature.csv', index=False)
    
    print(f"Feature-engineered training data saved to: {output_dir}train_feature.csv")
    print(f"Feature-engineered test data saved to: {output_dir}test_feature.csv")
    print(f"Final training set shape: {train_feature.shape}")
    print(f"Final test set shape: {test_feature.shape}")
    
    # 生成特征重要性报告
    if hasattr(feature_engineer, 'feature_selector') and feature_engineer.feature_selector is not None:
        feature_importance = feature_engineer.get_feature_importance()
        if feature_importance is not None:
            print("\nTop 20 most important features:")
            print(feature_importance.head(20))
            
            # 保存特征重要性
            feature_importance.to_csv('feature_importance.csv', index=False)
            print("Feature importance saved to: feature_importance.csv")
    
    # 生成特征描述文档
    generate_feature_description(train_feature, feature_engineer)
    
    return feature_engineer, train_feature, test_feature

def generate_feature_description(train_feature, feature_engineer):
    """
    生成特征描述文档
    """
    feature_info = []
    feature_info.append("# Ames住房数据集特征工程说明文档\n")
    feature_info.append("## 一、特征工程概述\n")
    feature_info.append("### 处理流程")
    feature_info.append("1. 分类特征编码 → 2. 新特征创建 → 3. 数值特征处理 → 4. 特征选择\n")
    feature_info.append(f"### 最终特征维度")
    feature_info.append(f"- 训练集: {train_feature.shape[1]} 个特征")
    feature_info.append(f"- 测试集: {train_feature.shape[1] - 1} 个特征 (不含目标变量)\n")
    
    feature_info.append("## 二、特征类别统计\n")
    
    # 分析特征类型
    numerical_features = []
    categorical_features = []
    engineered_features = []
    
    for col in train_feature.columns:
        if col == 'SalePrice' or col == 'Id':
            continue
        if '_encoded' in col:
            categorical_features.append(col)
        elif any(x in col for x in ['Total', 'Ratio', 'Age', 'Score', 'Is']):
            engineered_features.append(col)
        elif '_log' in col:
            numerical_features.append(col)
        else:
            numerical_features.append(col)
    
    feature_info.append(f"- 数值特征: {len(numerical_features)} 个")
    feature_info.append(f"- 分类特征(编码后): {len(categorical_features)} 个")
    feature_info.append(f"- 工程特征: {len(engineered_features)} 个")
    feature_info.append(f"- 总特征数: {len(numerical_features) + len(categorical_features) + len(engineered_features)} 个\n")
    
    feature_info.append("## 三、主要工程特征\n")
    
    # 新特征描述
    new_features_desc = {
        'TotalSF': '总建筑面积 = 地下室面积 + 一层面积 + 二层面积',
        'TotalArea': '总面积 = 总建筑面积 + 车库面积 + 木甲板面积 + 开放式门廊面积',
        'TotalBath': '总浴室当量 = 全卫 + 0.5*半卫 + 地下室全卫 + 0.5*地下室半卫',
        'TotalPorch': '总门廊面积 = 各类型门廊面积之和',
        'HouseAge': '房屋年龄 = 销售年份 - 建造年份',
        'RemodAge': '改造年龄 = 销售年份 - 改造年份',
        'IsRemodeled': '是否经过改造 (0/1)',
        'IsNewHouse': '是否新房 (0/1)',
        'ExteriorScore': '外部质量评分 = 外部质量 × 外部状况',
        'BasementScore': '地下室质量评分 = 地下室质量 × 地下室状况',
        'GarageScore': '车库质量评分 = 车库质量 × 车库状况'
    }
    
    for feature, desc in new_features_desc.items():
        if feature in train_feature.columns:
            feature_info.append(f"- **{feature}**: {desc}")
    
    feature_info.append("\n## 四、输出文件")
    feature_info.append("1. `feature_engineering.py` - 特征工程主程序")
    feature_info.append("2. `data/train_feature.csv` - 处理后的训练特征")
    feature_info.append("3. `data/test_feature.csv` - 处理后的测试特征")
    feature_info.append("4. `feature_importance.csv` - 特征重要性排名")
    feature_info.append("5. `feature_desc.txt` - 本说明文档")
    
    # 写入文件
    with open('feature_desc.txt', 'w', encoding='utf-8') as f:
        f.write('\n'.join(feature_info))
    
    print("Feature description saved to: feature_desc.txt")

if __name__ == "__main__":
    try:
        feature_engineer, train_feature, test_feature = main()
        print("\n" + "="*50)
        print("特征工程完成！")
        print("输出文件：")
        print("- data/train_feature.csv")
        print("- data/test_feature.csv") 
        print("- feature_importance.csv")
        print("- feature_desc.txt")
        print("="*50)
        
        # 等待用户按键后才关闭窗口
        input("按回车键退出...")
        
    except Exception as e:
        print(f"发生错误: {e}")
        import traceback
        traceback.print_exc()
        input("按回车键退出...")